#-*-coding:UTF-8-*-

__author__='zhaoxp'


import sys
import os
import os.path
import stat


def main():
    '''add id_rsa.pub content to ~/.ssh/authorized_keys
    check whether ~/.ssh/authorized_keys contains 
    '''
    if os.name != 'posix':
        sys.stderr.write('Error: OS is not posix and return\n')
        return
    if len(sys.argv) > 1:
        key_file_path = os.path.join('/tmp', sys.argv[1])
        key_file_content = ['']
        if not os.path.isfile(key_file_path):
            sys.stderr.write('key file not found\n')
            return
        with open(os.path.join(key_file_path)) as key_file:
            key_file_content = key_file.readlines()
        key_file_content = [l for l in key_file_content if l.strip()]
        #print 'key_file_content =', key_file_content

        line = ''
        if key_file_content:
            line = key_file_content[0]
        else:
            sys.stderr.write('key file empty\n')
            return
        username_plus_hostname = line.split(' ')[-1]

        if not username_plus_hostname:
            sys.stderr.write('key file content error\n')
            return

        ssh_path = os.path.join(os.environ['HOME'], '.ssh')
        if not os.path.isdir(ssh_path):
            os.mkdir(ssh_path)
            os.chmod(ssh_path, stat.S_IRUSR+stat.S_IWUSR+stat.S_IXUSR)

        #ak_file_path : file path of 'authorized_keys'
        ak_file_path = os.path.join(ssh_path, 'authorized_keys')

        '''
        1 authorized_keys existed
            1.1 key exists:
                1.1.1 same key               -> no action
                1.1.2 different key          -> update key
            1.2 key does not exist           -> append key
        2 not authorized_keys and create new -> create and add key
        '''
        if os.path.isfile(ak_file_path):
            ak_file_content = ['']
            with open(ak_file_path, 'r') as ak_file:
                ak_file_content = ak_file.readlines()
            ak_file_content = [l for l in ak_file_content if l.strip()]
            existed = False
            key_is_same = False
            for index, value in enumerate(ak_file_content):
                if value.find(username_plus_hostname) >= 0:
                    existed = True
                    if value == line:
                        key_is_same = True
                    else:
                        ak_file_content[index] = line
            if existed:
                if not key_is_same:
                    #update key
                    #print 'update key'
                    with open(ak_file_path, 'w') as ak_file:
                        for l in ak_file_content:
                            ak_file.write(l)
            else:
                #append key
                #print 'append key'
                with open(ak_file_path, 'a') as ak_file:
                    #ak_file.seek(0, 2)
                    ak_file.write(line)
        else:
            #create and add key
            #print 'create and add key'
            with open(ak_file_path, 'w') as ak_file:
                ak_file.write(line)
        sys.stdout.write('Successfull')
    else:
        sys.stderr.write('sys.argv length is less than 2\n')



if __name__=='__main__':
    main()